from sqlite import SQLiteDB
import os
from insightface.insightface import InsightFace

# 建立数据库
db = SQLiteDB('./db/sqlite/faces.db')

# 建立人脸向量表
face_vector_table_sql = '''
create table if not exists sqlite (id INTEGER PRIMARY KEY AUTOINCREMENT, name  TEXT,vector BLOB)
'''
db.execute(face_vector_table_sql)

# 加载模型，录入初始的人脸数据
model = InsightFace()
for file_name in os.listdir('./test_images'):
    img_path = os.path.join('./test_images', file_name)
    img, faces = model.load_and_process_image(img_path)
    face_feature = model.extract_face_features(faces)
    face_vector_insert_sql = '''
        insert into sqlite (name,vector) values (?,?)
    '''
    db.execute(face_vector_insert_sql, (file_name.split('.')[0], face_feature[0].tobytes()))

# 显示一下数据库的结果
rows = db.fetch('select * from sqlite')
print(f'插入总条数：{len(rows)}')
for row in rows:
    print(row)

db.close()
